package tx

import (
	"github.com/timandy/routine"
	"gorm.io/gorm"
)

var _tx = routine.NewInheritableThreadLocal[bool]()
var txRollback = routine.NewInheritableThreadLocal[bool]()
var txDb = routine.NewInheritableThreadLocal[*gorm.DB]()

func BeginTX() {
	_tx.Set(true)
}

func Clear() {
	_tx.Remove()
	txRollback.Remove()
	txDb.Remove()
}

func SetDB(db *gorm.DB) {
	txDb.Set(db)
}

func GetDB() *gorm.DB {
	return txDb.Get()
}

func NeedRollback() {
	txRollback.Set(true)
}

func TryCommit() {
	db := txDb.Get()
	if db != nil {
		if txRollback.Get() {
			db.Rollback()
		} else {
			db.Commit()
		}
	}
}

func IsInTX() bool {
	return _tx.Get()
}
